// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Public API for different memory sharing / IO mechanisms

#pragma once

#include <cstdint>
#include <memory>
#include <string_view>
#include <vector>

#include <nebula/io/concurrency.h>
#include <nebula/io/interfaces.h>
#include <nebula/types/type_fwd.h>
#include <turbo/base/macros.h>

namespace nebula {


    namespace io {

        /// \brief An output stream that writes to a resizable buffer
        class TURBO_EXPORT BufferOutputStream : public OutputStream {
        public:
            explicit BufferOutputStream(const std::shared_ptr<ResizableBuffer> &buffer);

            /// \brief create in-memory output stream with indicated capacity using a
            /// memory pool
            /// \param[in] initial_capacity the initial allocated internal capacity of
            /// the OutputStream
            /// \param[in,out] pool a MemoryPool to use for allocations
            /// \return the created stream
            static turbo::Result<std::shared_ptr<BufferOutputStream>> create(
                    int64_t initial_capacity = 4096, MemoryPool *pool = default_memory_pool());

            ~BufferOutputStream() override;

            // Implement the OutputStream interface

            /// close the stream, preserving the buffer (retrieve it with finish()).
            turbo::Status close() override;

            bool closed() const override;

            turbo::Result<int64_t> tell() const override;

            turbo::Status write(const void *data, int64_t nbytes) override;

            /// \cond FALSE
            using OutputStream::write;
            /// \endcond

            /// close the stream and return the buffer
            turbo::Result<std::shared_ptr<Buffer>> finish();

            /// \brief Initialize state of OutputStream with newly allocated memory and
            /// set position to 0
            /// \param[in] initial_capacity the starting allocated capacity
            /// \param[in,out] pool the memory pool to use for allocations
            /// \return turbo::Status
            turbo::Status reset(int64_t initial_capacity = 1024, MemoryPool *pool = default_memory_pool());

            int64_t capacity() const { return capacity_; }

        private:
            BufferOutputStream();

            // Ensures there is sufficient space available to write nbytes
            turbo::Status Reserve(int64_t nbytes);

            std::shared_ptr<ResizableBuffer> buffer_;
            bool is_open_;
            int64_t capacity_;
            int64_t position_;
            uint8_t *mutable_data_;
        };

        /// \brief A helper class to track the size of allocations
        ///
        /// Writes to this stream do not copy or retain any data, they just bump
        /// a size counter that can be later used to know exactly which data size
        /// needs to be allocated for actual writing.
        class TURBO_EXPORT MockOutputStream : public OutputStream {
        public:
            MockOutputStream() : extent_bytes_written_(0), is_open_(true) {}

            // Implement the OutputStream interface
            turbo::Status close() override;

            bool closed() const override;

            turbo::Result<int64_t> tell() const override;

            turbo::Status write(const void *data, int64_t nbytes) override;
            /// \cond FALSE
            using Writable::write;
            /// \endcond

            int64_t GetExtentBytesWritten() const { return extent_bytes_written_; }

        private:
            int64_t extent_bytes_written_;
            bool is_open_;
        };

        /// \brief An output stream that writes into a fixed-size mutable buffer
        class TURBO_EXPORT FixedSizeBufferWriter : public WritableFile {
        public:
            /// Input buffer must be mutable, will abort if not
            explicit FixedSizeBufferWriter(const std::shared_ptr<Buffer> &buffer);

            ~FixedSizeBufferWriter() override;

            turbo::Status close() override;

            bool closed() const override;

            turbo::Status Seek(int64_t position) override;

            turbo::Result<int64_t> tell() const override;

            turbo::Status write(const void *data, int64_t nbytes) override;
            /// \cond FALSE
            using Writable::write;
            /// \endcond

            turbo::Status write_at(int64_t position, const void *data, int64_t nbytes) override;

            void set_memcopy_threads(int num_threads);

            void set_memcopy_blocksize(int64_t blocksize);

            void set_memcopy_threshold(int64_t threshold);

        protected:
            class FixedSizeBufferWriterImpl;

            std::unique_ptr<FixedSizeBufferWriterImpl> impl_;
        };

        /// \class BufferReader
        /// \brief Random access zero-copy reads on an nebula::Buffer
        class TURBO_EXPORT BufferReader
                : public internal::RandomAccessFileConcurrencyWrapper<BufferReader> {
        public:
            /// \brief Instantiate from std::shared_ptr<Buffer>.
            ///
            /// This is a zero-copy constructor.
            explicit BufferReader(std::shared_ptr<Buffer> buffer);

            TURBO_DEPRECATED(
                    "Deprecated in 14.0.0. Use from_string or BufferReader(std::shared_ptr<Buffer> "
                    "buffer) instead.")
            explicit BufferReader(BufferSpan buffer);

            TURBO_DEPRECATED(
                    "Deprecated in 14.0.0. Use from_string or BufferReader(std::shared_ptr<Buffer> "
                    "buffer) instead.")
            BufferReader(const uint8_t *data, int64_t size);

            /// \brief Instantiate from std::string_view. Does not own data
            /// \deprecated Deprecated in 14.0.0. Use from_string or
            /// BufferReader(std::shared_ptr<Buffer> buffer) instead.
            TURBO_DEPRECATED(
                    "Deprecated in 14.0.0. Use from_string or BufferReader(std::shared_ptr<Buffer> "
                    "buffer) instead.")
            explicit BufferReader(std::string_view data);

            /// \brief Instantiate from std::string. Owns data.
            static std::unique_ptr<BufferReader> from_string(std::string data);

            bool closed() const override;

            bool supports_zero_copy() const override;

            std::shared_ptr<Buffer> buffer() const { return buffer_; }

            // Synchronous read_async override
            Future<std::shared_ptr<Buffer>> read_async(const IOContext &, int64_t position,
                                                      int64_t nbytes) override;

            turbo::Status will_need(const std::vector<ReadRange> &ranges) override;

        protected:
            friend RandomAccessFileConcurrencyWrapper<BufferReader>;

            turbo::Status DoClose();

            turbo::Result<int64_t> DoRead(int64_t nbytes, void *buffer);

            turbo::Result<std::shared_ptr<Buffer>> DoRead(int64_t nbytes);

            turbo::Result<int64_t> DoReadAt(int64_t position, int64_t nbytes, void *out);

            turbo::Result<std::shared_ptr<Buffer>> DoReadAt(int64_t position, int64_t nbytes);

            turbo::Result<std::string_view> DoPeek(int64_t nbytes) override;

            turbo::Result<int64_t> DoTell() const;

            turbo::Status DoSeek(int64_t position);

            turbo::Result<int64_t> DoGetSize();

            turbo::Status CheckClosed() const {
                if (!is_open_) {
                    return turbo::invalid_argument_error("Operation forbidden on closed BufferReader");
                }
                return turbo::OkStatus();
            }

            std::shared_ptr<Buffer> buffer_;
            const uint8_t *data_;
            int64_t size_;
            int64_t position_;
            bool is_open_;
        };

    }  // namespace io
}  // namespace nebula
